from io import BytesIO
from PIL import Image
from common import errorCode
from flask import Flask, request
from common.customOutput import custom_response
from controllers.chinese_ocr_lite import OcrHandle
from conf.params_config import *
from handle.start import post_process
from common.stringUtils import LevenshteinRatio
import json
from yxlog.logger import get_logger
import traceback


ocrHandle = OcrHandle()
logger = get_logger()


app = Flask(__name__)


@app.route("/api/v1.1/check", methods=["GET"])
def api_service_check_handler():
    result = "API Service Is Available"
    return custom_response(data=result)


@app.route("/api/v1.1/certificateOCR", methods=["POST"])
def certificate_ocr_handler():
    image = request.files.get("srcImage")
    src_data = request.form.get("srcData")
    if not image or not src_data:
        logger.error("上传文件字段缺失")
        return custom_response(code=errorCode.ErrorParameter)

    image = image.read()
    try:
        image = Image.open(BytesIO(image)).convert('RGB')
    except Exception as e:
        logger.error(e)
        return custom_response(code=errorCode.ErrorParser)

    try:
        src_data = json.loads(src_data)
    except Exception as e:
        logger.error(e)
        return custom_response(code=errorCode.ErrorParser)
    try:
        short_size = 960
        img_w, img_h = image.size
        if max(img_w, img_h) * (short_size * 1.0 / min(img_w, img_h)) > max_longer_size:
            short_size = 480

        result, angle_res = ocrHandle.text_predict(image, short_size)
        if angle_res:
            for rotate in [Image.ROTATE_270, Image.ROTATE_90, Image.ROTATE_180]:
                image = image.transpose(rotate)
                result, angle_res = ocrHandle.text_predict(image, short_size)
                if not angle_res:
                    break
        result = post_process(result, image)
    except Exception as e:
        traceback.print_exc()
        logger.error(e)
        return custom_response(code=errorCode.ErrorInnerService)

    if not result:
        return custom_response(code=errorCode.FailedParser)
    if result["certificateName"] == "UNKNOWN":
        return custom_response(data=result, code=errorCode.FailedMapping)
    result = LevenshteinRatio(result, src_data).result

    return custom_response(data=result)


@app.route("/api/v1.1/textDetection", methods=["POST"])
def certificate_detection_handler():
    image = request.files.get("srcImage")
    if not image:
        logger.error("上传文件字段缺失")
        return custom_response(code=errorCode.ErrorParameter)

    image = image.read()
    try:
        image = Image.open(BytesIO(image)).convert('RGB')
    except Exception as e:
        logger.error(e)
        return custom_response(code=errorCode.ErrorParser)

    try:
        short_size = 960
        img_w, img_h = image.size
        if max(img_w, img_h) * (short_size * 1.0 / min(img_w, img_h)) > max_longer_size:
            short_size = 480

        result, angle_res = ocrHandle.text_predict(image, short_size)
        if angle_res:
            for rotate in [Image.ROTATE_270, Image.ROTATE_90, Image.ROTATE_180]:
                image = image.transpose(rotate)
                result, angle_res = ocrHandle.text_predict(image, short_size)
                if not angle_res:
                    break
    except Exception as e:
        traceback.print_exc()
        logger.error(e)
        return custom_response(code=errorCode.ErrorInnerService)

    return custom_response(data=result)


@app.route("/api/v1.1/textRecognize", methods=["POST"])
def certificate_recognize_handler():
    image = request.files.get("srcImage")
    if not image:
        logger.error("上传文件字段缺失")
        return custom_response(code=errorCode.ErrorParameter)

    image = image.read()
    try:
        image = Image.open(BytesIO(image)).convert('RGB')
    except Exception as e:
        logger.error(e)
        return custom_response(code=errorCode.ErrorParser)

    result = ocrHandle.text_rec(image)
    return custom_response(data=result)


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8000, debug=False)
